﻿#light

namespace CodeBlock.AMD64

open System

open CodeBlock

open CodeBlock.Utils

type Register =
    | AX = 0
    | CX = 1
    | DX = 2
    | BX = 3
    | SP = 4
    | BP = 5
    | SI = 6
    | DI = 7
    | R8 = 8
    | R9 = 9
    | RA = 10
    | RB = 11
    | RC = 12
    | RD = 13
    | RE = 14
    | RF = 15

type Condition =
    // CF = 1
    | Carry = 0
    // CF = 1
    | NotAboveOrEqual = 0
    // CF = 0
    | NotBelow = 1
    // CF = 0
    | NotCarry = 1
    // CF = 0
    | AboveOrEqual = 1
    // ZF = 0
    | Zero = 2
    // ZF = 0
    | Equal = 2
    // ZF = 1
    | NotEqual = 3
    // ZF = 1
    | NotZero = 3
    // CF = 1 or ZF = 1
    | BelowOrEqual = 4
    // CF = 1 or ZF = 1
    | NotAbove = 4
    // CF = 0 and ZF = 0
    | NotBelowOrEqual = 5
    // CF = 0 and ZF = 0
    | Above = 5
    // SF = 1
    | Sign = 6
    // SF = 0
    | NotSign = 7
    // PF = 1
    | Parity = 8
    // PF = 1
    | ParityEven = 8
    // PF = 0
    | NotParity = 9
    // PF = 0
    | ParityOdd = 9
    // SF <> OF
    | Less = 10
    // SF <> OF
    | NotGreaterOrEqual = 10
    // SF = OF
    | NotLess = 11
    // SF = OF
    | GreaterOrEqual = 11
    // ZF = 1 or SF <> OF
    | LessOrEqual = 12
    // ZF = 1 or SF <> OF
    | NotGreater = 12
    // ZF = 0 and SF = OF
    | NotLessOrEqual = 13
    // ZF = 0 and SF = OF
    | Greater = 13


type RegisterOperand(reg: Register, size: int)=
    member this.Register with get() = reg
    member this.Size with get() = size
    
    override this.ToString() =
        let prefix =
            match size with
            | 2 -> ""
            | 4 -> "E"
            | 8 -> "R"
            | 1 -> "_"
            | _ -> "S" + size.ToString()
        
        prefix + (reg.ToString().ToUpper())

type ImmediateOperand(?baseSymbol: string, ?value: int64) =
    let baseSymbol = make_value baseSymbol null
    let value = make_value value 0L
    
    member this.Base with get() = baseSymbol
    member this.Value with get() = value
    
    interface IEquatable<ImmediateOperand> with
        override this.Equals(other: ImmediateOperand)=
            (baseSymbol = other.Base) && (value = other.Value)
    
    override this.ToString() =
        match baseSymbol, value with
        | null, _ -> "~*" + value.ToString()
        | _, 0L -> "\"" + baseSymbol + "\""
        | _, _ -> "(\"" + baseSymbol + "\" + " + value.ToString() + ")"
    
    static member (+) (left: ImmediateOperand, right: ImmediateOperand) =
        let baseSymbol =
            if left.Base = null then right.Base
            elif right.Base = null then left.Base
            else raise(InvalidOperationException("can't add immediate operands both having bases"))
        ImmediateOperand(baseSymbol = baseSymbol, value = left.Value + right.Value)
    
    static member Zero with get() = ImmediateOperand()
    
    [<OverloadID("imm_int")>]
    static member (+) (left: ImmediateOperand, right)=
        ImmediateOperand(baseSymbol = left.Base, value = left.Value + right)
    
    [<OverloadID("imm_int64")>]
    static member (+) (left: ImmediateOperand, right)=
        left + (int64 right)

type MemoryOperand(?original: MemoryOperand, ?displacement: ImmediateOperand, ?baseRegister: Register, ?indexRegister: Register, ?scale: int, ?rip: bool)=
    let rip = make_value rip (if original = None then false else original.Value.UseRIP)
    let displacement = make_value displacement (if original = None then ImmediateOperand.Zero else original.Value.Displacement)
    let scale = make_value scale (if original = None then 1 else original.Value.Scale)
    let baseRegister =
        if baseRegister = None then
            if original = None then None else original.Value.Base
        else
            baseRegister
    let indexRegister =
        if indexRegister = None then
            if original = None then None else original.Value.Index
        else
            indexRegister
    
    member this.Displacement with get() = displacement
    member this.Base with get() = baseRegister
    member this.Index with get() = indexRegister
    member this.Scale with get() = scale
    member this.UseRIP with get() = rip
    
    override this.ToString() =
        if rip then
            "(RIP + " + displacement.Value.ToString() + ")"
        else
            let baseRegister = if baseRegister = None then null else "R" + baseRegister.Value.ToString()
            let index = if indexRegister = None then null else "R" + indexRegister.Value.ToString() + "*" + scale.ToString() 
            if (baseRegister = null) && (index = null) then
                "(MEM (" + displacement.ToString() + "))"
            elif baseRegister = null then
                "("+index + " + " + displacement.ToString()+")"
            elif index = null then
                "("+baseRegister + " + " + displacement.ToString()+")"
            else
                "("+baseRegister + " + " + index + " + " + displacement.ToString()+")"

type Operand = 
    | Immediate of ImmediateOperand
    | Register of RegisterOperand
    | Memory of MemoryOperand
    
    override this.ToString() =
        match this with
        | Immediate(imm) -> imm.ToString()
        | Register(reg) -> reg.ToString()
        | Memory(mem) -> mem.ToString()
    
    static member (*) (reg, scale) =
        match reg with
        | Register(reg) ->
            Memory(MemoryOperand(indexRegister = reg.Register, scale = scale))
        | _ ->
            raise(InvalidOperationException("can't multiply non-register operand by integer value"))
    
    [<OverloadID("op_op")>]
    static member (+) (left, right) =
        match left, right with
        | Memory(left), Immediate(right) when left.UseRIP && (right.Base = null) ->
            Memory(MemoryOperand(displacement = left.Displacement + right, rip = true))
        | Memory(mem), Register(reg) | Register(reg), Memory(mem) ->
            if mem.UseRIP then
                raise(InvalidOperationException("can't add register to RIP-relative address"))
            match mem.Base, mem.Index with
            | None, None ->
                Memory(MemoryOperand(displacement = mem.Displacement, baseRegister = reg.Register, scale = mem.Scale))
            | None, Some(register) ->
                Memory(MemoryOperand(displacement = mem.Displacement, baseRegister = reg.Register, indexRegister = register, scale = mem.Scale))
            | Some(register), None ->
                Memory(MemoryOperand(displacement = mem.Displacement, baseRegister = reg.Register, indexRegister = register, scale = 1))
            | _, _ ->
                raise(InvalidOperationException("can't add another register"))
        | Register(left), Register(right) ->
            Memory(MemoryOperand(baseRegister = left.Register, indexRegister = right.Register, scale = 1))
        | Register(reg), Immediate(imm) | Immediate(imm), Register(reg) ->
            Memory(MemoryOperand(displacement = imm, baseRegister = reg.Register))
        | _, _ ->
            raise(InvalidOperationException("can't add two operands"))
    
    [<OverloadID("op_int")>]
    static member (+) (left, right: int) =
        match left with
        | Memory(left) ->
            Memory(MemoryOperand(original = left, displacement = left.Displacement + right))
        | Register(left) ->
            Memory(MemoryOperand(displacement = ImmediateOperand.Zero + right, baseRegister = left.Register))
        | Immediate(imm) ->
            Memory(MemoryOperand(displacement = imm + right))

type Operation=
    | Comment   = -1
    | Add       = 0
    | Adc       = 1
    | And       = 2
    | Sub       = 3
    | Sbb       = 4
    | Xor       = 5
    | Or        = 6
    | Je        = 7
    | Jae       = 8
    | Jge       = 9
    | Jbe       = 10
    | Jle       = 11
    | Jgt       = 12
    | Jlt       = 13
    | Ja        = 14
    | Jb        = 15
    | Jne       = 16
    | Jmp       = 17
    | Jz        = 18
    | Jnz       = 19
    | Call      = 20
    | Movsx     = 21
    | Movsxd    = 22
    | Cdq       = 23
    | Cqo       = 24
    | Mov       = 25
    | Div       = 26
    | Idiv      = 27
    | Mul       = 28
    | Imul      = 29
    | Push      = 30
    | Pop       = 31
    | Lea       = 32
    | Neg       = 33
    | Nop       = 34
    | Not       = 35
    | Ret       = 36
    | Shl       = 37
    // unsigned shift right
    | Shr       = 38
    // signed shift right
    | Sar       = 39
    | Label     = 40
    // lock prefix
    | Lock      = 41
    // rep prefix
    | Rep       = 42
    // set direction flag
    | Std       = 43
    | Movsq     = 44
    | Movsb     = 45

type InstructionAMD64 =
    {   Operation: Operation;
        Operands: Operand list; }
    
    override this.ToString() =
        let rec list_items items =
            match items with
            | [] -> ""
            | head :: others ->
                head.ToString() + " " + list_items others
        
        let operands = list_items this.Operands
        Enum.GetName(typeof<Operation>, box this.Operation).ToUpper() + " " + operands

module internal Assembly = 
    let public RAX = Register(RegisterOperand(Register.AX, 8))
    let public RBX = Register(RegisterOperand(Register.BX, 8))
    let public RCX = Register(RegisterOperand(Register.CX, 8))
    let public RDX = Register(RegisterOperand(Register.DX, 8))
    let public RBP = Register(RegisterOperand(Register.BP, 8))
    let public RSP = Register(RegisterOperand(Register.SP, 8))
    let public RSI = Register(RegisterOperand(Register.SI, 8))
    let public RDI = Register(RegisterOperand(Register.DI, 8))
    
    let public EAX = Register(RegisterOperand(Register.AX, 4))
    let public EBX = Register(RegisterOperand(Register.BX, 4))
    let public ECX = Register(RegisterOperand(Register.CX, 4))
    let public EDX = Register(RegisterOperand(Register.DX, 4))
    
    let public RIP = Memory(MemoryOperand(rip = true))
    
    // builds immediate operand
    let inline public ( ~* ) value = Immediate(ImmediateOperand(value = int64 value))
    // builds immediate operand
    let public (~&) value = Immediate(ImmediateOperand(baseSymbol = value))
    
    let public MEM value = Memory(MemoryOperand(displacement = value))
    
    let public BINARY op dest source (code: InstructionAMD64 list) =
        { Operation = op; Operands = dest :: source :: []; } :: code
    
    let public UNARY op target code = 
        { Operation = op; Operands = target :: []; } :: code
    
    let public SIMPLE op code =
        { Operation = op; Operands = []; } :: code
    
    let public ADD dest source =
        BINARY Operation.Add dest source
    
    let public ADC dest source =
        BINARY Operation.Adc dest source
    
    let public AND dest source =
        BINARY Operation.And dest source
    
    let public SUB dest source =
        BINARY Operation.Sub dest source
    
    let public SBB dest source =
        BINARY Operation.Sbb dest source
    
    let public XOR dest source =
        BINARY Operation.Xor dest source
    
    let public OR dest source =
        BINARY Operation.Or dest source
    
    let public JE target =
        UNARY Operation.Je target
    
    let public JAE target =
        UNARY Operation.Jae target
    
    let public JGE target =
        UNARY Operation.Jge target
    
    let public JBE target =
        UNARY Operation.Jbe target
    
    let public JLE target =
        UNARY Operation.Jle target
    
    let public JGT target =
        UNARY Operation.Jgt target
    
    let public JLT target =
        UNARY Operation.Jlt target
    
    let public JA target =
        UNARY Operation.Ja target
    
    let public JB target =
        UNARY Operation.Jb target
    
    let public JNE target =
        UNARY Operation.Jne target
    
    let public JMP target =
        UNARY Operation.Jmp target
    
    let public JZ target =
        UNARY Operation.Jz target
    
    let public JNZ target =
        UNARY Operation.Jnz target
    
    let public CALL target =
        UNARY Operation.Call target
    
    let public MOVSX dest source =
        BINARY Operation.Movsx dest source
    
    let public MOVSXD dest source =
        BINARY Operation.Movsxd dest source
    
    let public MOVSQ =
        SIMPLE Operation.Movsq
    
    let public MOVSB =
        SIMPLE Operation.Movsb
    
    let public MOV dest source =
        BINARY Operation.Mov dest source
    
    let public DIV divisor = 
        UNARY Operation.Div divisor
    
    let public IDIV divisor = 
        UNARY Operation.Idiv divisor
    
    let public MUL multiplier =
        UNARY Operation.Mul multiplier
    
    let public IMUL multiplier =
        UNARY Operation.Imul multiplier
    
    let public PUSH value =
        UNARY Operation.Push value
    
    let public POP target =
        UNARY Operation.Pop target
    
    let public LEA dest source =
        BINARY Operation.Lea dest source
    
    let public NEG target =
        UNARY Operation.Neg target
    
    let public NOP =
        SIMPLE Operation.Nop
    
    let public REP =
        SIMPLE Operation.Rep
    
    let public STD =
        SIMPLE Operation.Std
    
    let public CDQ =
        SIMPLE Operation.Cdq
    
    let public CQO =
        SIMPLE Operation.Cqo
    
    let public NOT target = 
        UNARY Operation.Not target
    
    let public RET bytes =
        UNARY Operation.Ret bytes
    
    let public SHL dest shift = 
        BINARY Operation.Shl dest shift
    
    let public SHR dest shift =
        BINARY Operation.Shr dest shift
    
    let public SAR dest shift =
        BINARY Operation.Sar dest shift
    
    let public LABEL name =
        UNARY Operation.Label name
    
    let public COMMENT str =
        UNARY Operation.Comment (Immediate(ImmediateOperand(baseSymbol = str)))
    
    let public ( ~/ ) (name: string) = LABEL (Immediate(ImmediateOperand(baseSymbol = name)))
    
    // Reverses instruction list.
    // This function was added to speedup assembly generation
    let public COMMIT code =
        List.rev code
    
//    let VADD =
//        []
//        |> POP RAX
//        |> POP RBX
//        |> ADD RAX RBX
//        |> PUSH RAX
//        |> COMMIT
//    
//    let mymegacode = 
//        []
//        |> ADD RAX ~*100
//        |> PUSH RAX
//        |> COMMIT

[<AbstractClass>]
type AbstractAssembler()=
    abstract Assemble: InstructionAMD64 list -> CodeBlock



//type DestSourceInstruction(dest, source, prefixes)=
//    inherit Instruction(dest::source::[], prefixes)
//    member this.Dest with get() = this.Operands
//    member this.Source with get() = source
//
//type AluInstruction(dest, source, prefixes) =
//    inherit DestSourceInstruction(dest, source, prefixes)
//
//type ImplicitInstruction() =
//    inherit Instruction([],[])
//
//type Add(dest, source, prefixes)=
//    inherit AluInstruction(dest, source, prefixes)
//
//type Adc(dest, source, prefixes) =
//    inherit AluInstruction(dest, source, prefixes)
//
//type And(dest, source, prefixes) =
//    inherit AluInstruction(dest, source, prefixes)
//
////Sign-extend AL into AX
//type Cbw()=
//    inherit ImplicitInstruction()
//
////Sign-extend AX into EAX
//type Cwde()=
//    inherit ImplicitInstruction()
//
////Sign-extend EAX into RAX
//type Cdqe()=
//    inherit ImplicitInstruction()
//
////Sign-extend AX into DX:AX
//type Cwd()=
//    inherit ImplicitInstruction()
//
////Sign-extend EAX into EDX:EAX
//type Cdq()=
//    inherit ImplicitInstruction()
//
////Sign-extend RAX into RDX:RAX
//type Cqo()=
//    inherit ImplicitInstruction()
//
//type Cmp(dest, source, prefixes)=
//    inherit AluInstruction(dest, source, prefixes)
//
//type Cpuid()=
//    inherit ImplicitInstruction()
//
//type Div(divisor)=
//    inherit Instruction(divisor::[],[])
//
//type Idiv(divisor)=
//    inherit Instruction(divisor::[],[])
//
//type Imul(multiplier)=
//    inherit Instruction(multiplier::[],[])
//
//type Int(vector)=
//    inherit Instruction(vector::[],[])
//
//type Jcc(target, condition: Condition)=
//    inherit Instruction(target::[], [])
//
//type JmpNear(target)=
//    inherit Instruction(target::[], [])
//
//type Lea(dest, memory)=
//    inherit DestSourceInstruction(dest, memory, [])
//
//type Mov(dest, source)=
//    inherit DestSourceInstruction(dest, source, [])
//
//type Movsx(dest, source)=
//    inherit DestSourceInstruction(dest, source, [])
//
//type Movsxd(dest, source)=
//    inherit DestSourceInstruction(dest, source, [])
//
//type Movzx(dest, source)=
//    inherit DestSourceInstruction(dest, source, [])
//
//type Mul(multiplier)=
//    inherit Instruction(multiplier::[], [])
//
//type Neg(dest)=
//    inherit Instruction(dest::[], [])
//
//type Nop()=
//    inherit Instruction([], [])
//
//type Not(dest)=
//    inherit Instruction(dest::[],[])
//
//type Or(dest, source, prefixes)=
//    inherit AluInstruction(dest, source, prefixes)
//
//type Pop(dest)=
//    inherit Instruction(dest::[], [])
//
//type Push(source)=
//    inherit Instruction(source::[], [])
//
//type Ret(imm)=
//    inherit Instruction(imm::[], [])
//
//type Shl(dest, source)=
//    inherit DestSourceInstruction(dest, source, [])
//
//
